# 运行过程中要对原始图片鼠标左键点选四个角点，顺序必须为左上、右上、左下，右下
# https://blog.csdn.net/weixin_42326479/article/details/102938733
import cv2
import os
import numpy as np
import tkinter as tk
from tkinter import filedialog

np.set_printoptions(suppress=True)

root = tk.Tk()
root.geometry("500x500")
root.title("图片矫正（Z字形选中四个点）")

# from tqdm import tqdm

# arr = [[0,0],[img.shape[1],0],[0,img.shape[0]],[img.shape[1],img.shape[0]]]
arr = []
# rc_FILE = "D:/data/IMG_20230719_105053.jpg"
img = None


def on_EVENT_LBUTTONDOWN(event, x, y, flags, param):
    if event == cv2.EVENT_LBUTTONDOWN:
        arr.append([x, y])
        cv2.circle(img, (x, y), 8, (255, 0, 255), thickness=3)
        cv2.imshow("manual", img)


def doCV2(src_FILE):
    global arr
    arr = []
    global img
    img = cv2.imread(src_FILE)  #中文目录会读不到
    #img = cv2.imdecode(np.fromfile(src_FILE, dtype=np.uint8), -1)
    img_transform = img.copy()

    cv2.namedWindow("manual", cv2.WINDOW_FREERATIO)
    cv2.setMouseCallback("manual", on_EVENT_LBUTTONDOWN)
    cv2.imshow("manual", img)
    # Esc 销毁窗口
    if len(arr) == 4 or cv2.waitKey(0) == 27:
        cv2.destroyAllWindows()
        if len(arr) == 3:
            print("选定的三个角点为：{}{}{}".format(arr[0], arr[1], arr[2]))
            return
        elif len(arr) == 2:
            print("选定的两个角点为：{}{}".format(arr[0], arr[1]))
            return
        elif len(arr) == 1:
            print("选定的一个角点为：{}".format(arr[0]))
            return
    # if len(arr) == 2:   #根据平行关系，自动补充p3，p4
    #     arr.append([arr[0][0], img.shape[0] - arr[0][1]])  # (p1.x,h-p2.y)
    #     arr.append([arr[0][1], img.shape[0]])  # (p2.x,h)
    P1, P2, P3, P4 = arr[:]

    print("选定的四个角点为：{}{}{}{}".format(P1, P2, P3, P4))  # Z字型选中四点
    cv2.circle(img_transform, tuple(P1), 2, (255, 0, 0), 3)
    cv2.circle(img_transform, tuple(P2), 2, (255, 0, 0), 3)
    cv2.circle(img_transform, tuple(P3), 2, (255, 0, 0), 3)
    cv2.circle(img_transform, tuple(P4), 2, (255, 0, 0), 3)
    # # 输入梯形的四个顶点
    srcPoints = np.vstack((P1, P2, P3, P4))
    srcPoints = np.float32(srcPoints)
    # 目标的像素值大小
    long = img.shape[1]  # 图片的长宽
    short = img.shape[0]
    # 设置目标画布的大小
    canvasPoints = np.array([[0, 0], [int(long), 0], [0, int(short)], [int(long), int(short)]])
    canvasPoints = np.float32(canvasPoints)
    # 计算转换矩阵
    perspectiveMatrix = cv2.getPerspectiveTransform(srcPoints, canvasPoints)
    print(perspectiveMatrix)
    # 完成透视变换
    perspectiveImg = cv2.warpPerspective(img_transform, perspectiveMatrix, (int(long), int(short)))
    print("展示校正结果")
    cv2.namedWindow(os.path.basename(src_FILE), cv2.WINDOW_FREERATIO)
    cv2.imshow(os.path.basename(src_FILE), perspectiveImg)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    dir_out = os.path.join(os.path.dirname(src_FILE), "cr")
    if not os.path.exists(dir_out):
        os.mkdir(dir_out)
    dst_FILE = os.path.join(dir_out, os.path.basename(src_FILE))  # "D:/data/real/out/result.JPG"
    cv2.imwrite(dst_FILE, perspectiveImg)  # 中文目录写入不成功，需要手动保存
    print("保存结果完毕")


def choose_image():
    src_FILE = filedialog.askopenfilename(initialdir="images/", title="选择图片",
                                          filetypes=(("JPEG files", "*.jpg"), ("PNG files", "*.png")))
    print(src_FILE)
    doCV2(src_FILE)


button = tk.Button(root, text="选择图片", command=choose_image)
button.pack(side="top", pady=50)
root.mainloop()
